import cv2
import time
import argparse
from picamera2 import Picamera2


def camera_test():
    """
    test camera with video
    """
    # 初始化picamera2
    picam2 = Picamera2()
    config = picam2.create_video_configuration(
        main={"size": (640, 480)},
        controls={"FrameRate": 30}
    )
    picam2.configure(config)
    picam2.start()
    number = 0
    # 读取帧（转为OpenCV格式）
    while True:
        frame = picam2.capture_array()
        # frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        cv2.imshow("CSI Camera", frame)
        print("Frame number:%d, time: %d", number, time.time())
        number += 1
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cv2.destroyAllWindows()


def video_record_by_frame():
    # 初始化picamera2
    picam2 = Picamera2()
    config = picam2.create_still_configuration(
        main={"size": (640, 480)},
        controls={"FrameRate": 30}
    )
    picam2.configure(config)
    picam2.start()

    # 设置视频编码格式
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter('output.avi', fourcc, 20.0, (640, 480))

    while True:
        # 读取一帧图像
        frame = cv2.cvtColor(picam2.capture_array(), cv2.COLOR_RGB2BGR)
        # 显示图像
        cv2.imshow('frame', frame)

        # 写入视频文件
        out.write(frame)

        # 按 'q' 键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放资源
    out.release()
    cv2.destroyAllWindows()

def picture_record_by_click():
    """
    once click, get and save a frame
    frame name is timestamp
    """
    # 初始化picamera2
    picam2 = Picamera2()
    config = picam2.create_still_configuration(
        main={"size": (640, 480)},
        controls={"FrameRate": 30}
    )
    picam2.configure(config)
    picam2.start()
    
    print("Press 's' to save a picture, 'q' to quit.")

    while True:
        # 读取一帧图像
        frame = cv2.cvtColor(picam2.capture_array(), cv2.COLOR_RGB2BGR)

        # 显示图像
        cv2.imshow('frame', frame)

        # 按 's' 键保存图片，按 'q' 键退出
        key = cv2.waitKey(1) & 0xFF
        if key == ord('s'):
            timestamp = time.strftime("%m%d_%H%M%S")
            cv2.imwrite(f'{timestamp}.jpg', frame)
            print(f"Saved {timestamp}.jpg")
        elif key == ord('q'):
            break

    # 释放资源
    cv2.destroyAllWindows()

def main():
    """
    main function
    """
    parser = argparse.ArgumentParser(description="camera test")
    parser.add_argument("--task", type=str, default="video_record_by_frame", help="task name")
    args = parser.parse_args()
    task_name = args.task

    if task_name == "video_record_by_frame":
        video_record_by_frame()
    elif task_name == "picture_record_by_click":
        picture_record_by_click()
    elif task_name == "camera_test":
        camera_test()
    else:
        print("task name error")


if __name__ == "__main__":
    main()